#include "ads.h"
#include "jobqueue.h"
#include "file.h"
#include "thing.h"

namespace Ads OPENBRACE

// For debugging. 
string Tuple::ToString() const { 
  ostringstream ostr;
  ostr << weight_ << " : " << "( " ;
  for (uint64 i=0; i<WIDTH; i++) ostr << a_[i] << " ";
  ostr << ")";
  return ostr.str();
}
 
// for debugging - reads the weight then the tuple from an input stream
bool Tuple::Read(istream &input) {
  if (!(input >> weight_)) return false;
  for (uint64 i=0; i<WIDTH; i++) 
    if (!(input >> a_[i])) return false;
  return true;
}

// for debugging - reads the length then the prefix from an input stream
bool Tuple::ReadPrefix(istream &input) {
  if (!(input >> prefix_length__)) return false;
  if (!(prefix_length__ < WIDTH)) {
    cout << "prefix_length__ = " << prefix_length__ 
	 << "  WIDTH = " << WIDTH << endl;
    CHECK(false);
  }
  for (uint64 i=0; i<prefix_length__; i++)
    if (!(input >> a_[i])) return false;
  return true;
}

// Create a prefix from variable args.
Tuple Prefix(uint64 prefix_length, ...) {
  CHECK(prefix_length < WIDTH);
  va_list va;
  va_start(va, prefix_length);
  Tuple r;
  r.prefix_length__ = prefix_length;
  for (uint64 i=0; i<prefix_length; i++) r.a_[i] = va_arg(va, uint64);	
  va_end(va);
  for (uint64 i=prefix_length; i<WIDTH; i++) r.a_[i] = 0;
  return r;
}

// create an entire tuple
Tuple FullTuple(uint64 weight, ...) {
  va_list va;
  va_start(va, weight);
  Tuple r;
  r.weight_ = weight;
  for (uint64 i=0; i<WIDTH; i++) r.a_[i] = va_arg(va, uint64);	
  va_end(va);
  return r;  
}


// The set of tuples with a particular prefix is [lower_bound upper_bound) .
// Here we compute those bounds.
Tuple Tuple::PrefixToLowerBound() const {
  Tuple ret;
  CHECK(prefix_length__ <= WIDTH);
  for (uint64 i=0; i<prefix_length__; i++) ret.a_[i] = a_[i];
  for (uint64 i=prefix_length__; i<WIDTH; i++) ret.a_[i] = 0;
  return ret;
}
Tuple Tuple::PrefixToUpperBound() const {
  Tuple ret;
  if (prefix_length__==0) { 
    // minor hack.  This works unless there is a tuple consisting entirely
    // of MAX_UINT64.  
    for (uint64 i=0; i<WIDTH; i++) ret.a_[i] = ~0ull;
    return ret;
  }
  for (uint64 i=0; i+1<prefix_length__; i++) ret.a_[i] = a_[i];
  ret.a_[prefix_length__-1] = a_[prefix_length__-1] + 1;
  for (uint64 i=prefix_length__; i<WIDTH; i++) ret.a_[i] = 0;
  return ret;
}

uint64 Tuple::NodeCount() const {
  uint64 node_count = 0;
  for (uint64 i=0; i<WIDTH; i++) 
    node_count += Thing(a_[i]).NodeCount();
  return node_count;
}

uint64 Tuple::WeightPerNode() const {
  return weight_ / NodeCount();
}

string Sampler::ToString() {
  ostringstream ostr;
  ostr << "Sampler size=" << size_ << " capacity=" << capacity_ << " {"
       << endl;
  for (uint64 i=0; i<size_; i++) {
    ostr << data_[i].ToString() << endl;
  }
  ostr << "}" << endl;
  return ostr.str();
}
void Sampler::Allocate(uint64 capacity) {
  size_ = 0;
  capacity_ = capacity;
  delete [] data_;
  data_ = new Tuple[capacity_];
  built_ = false;
}
  /*
void Sampler::ReadData(CircularBuffer * buffer) {
  CHECK(!built_);
  char * cp = (char *)(data_ + size_);
  buffer->Read(&cp, (capacity_ - size_) * sizeof(Tuple));
  size_ = ((Tuple *)cp) - data_;
}
  */
void Sampler::AddTuple(const Tuple & r) {
  CHECK(!built_);
  data_[size_] = r;
  size_++;
}
Sampler::~Sampler() { delete [] data_;}
void Sampler::SumWeights() {
  for (uint64 i=1; i<size_; i++) 
    data_[i].weight_ += data_[i-1].weight_;
}
void Sampler::UnSumWeights() {
  for (int64 i=size_-1; i>0; i--) 
    data_[i].weight_ -= data_[i-1].weight_;  
}

void Sampler::Verify() { 
  return;
  uint64 prev = 0;
  for (uint64 i=0; i<=size_; i++) {
    uint64 next = WeightBefore(data_ + i);
    CHECK(next >= prev);
    prev = next;
  }
}

void Sampler::Build() {
  CHECK(!built_);
  sort(data_, data_ + size_, less<Tuple>()); // could use JQSort here. 
  SumWeights();
  built_ = true;
  Verify();
}

void Sampler::Store(string filename) {
  File::Write(filename, StoredMemblock());
}

// new the sampler then call this
void Sampler::Load(string filename) {
  uint64 file_byte_size = File::Size(filename);
  uint64 tuple_size = file_byte_size / sizeof(Tuple);
  Allocate(tuple_size);
  File::Read(filename, Memblock((char*) data_, file_byte_size));
  size_ = tuple_size;
  built_ = true;
  Verify();
}

Sampler * Sampler::Merge(Sampler *as1, Sampler *as2) {  
  // TODO: merge weights correctly.
  
  uint64 pos1 = 0;
  uint64 pos2 = 0;
  uint64 size1 = as1->size_;
  uint64 size2 = as2->size_;
  uint64 weight1 = 0;
  uint64 weight2 = 0;
  Sampler * as = new Sampler;
  as->Allocate(size1 + size2);
  for (uint64 rep = 0; rep < size1 + size2; rep++) {
    if ( (pos2 == size2) || 
	 ((pos1 < size1) && (as1->data_[pos1] < as2->data_[pos2]))) {
      weight1 = as1->data_[pos1].weight_;      
      as->AddTuple(as1->data_[pos1]);
      pos1++;
    } else {
      weight2 = as2->data_[pos2].weight_;
      as->AddTuple(as2->data_[pos2]);
      pos2++;
    }
    as->data_[as->size_-1].weight_ = weight1 + weight2;
  }
  as->built_ = true;
  as->Verify();
  CHECK(as->size_ == size1 + size2);
  return as;
}

uint64 Sampler::WeightBefore(Tuple *p) {
  if (p == data_) return 0;
  return p[-1].weight_;
}

// sample a tuple in the range [begin end)
Tuple * Sampler::Sample(const Tuple & begin, const Tuple & end, 
			RNG *rng, SampleStrategy strategy) {
  CHECK(built_);
  Tuple * begin_p = lower_bound(data_, data_ + size_, begin);
  Tuple * end_p = lower_bound(data_, data_ + size_, end);
  uint64 begin_w = WeightBefore(begin_p);
  uint64 end_w = WeightBefore(end_p);
  CHECK(end_w >= begin_w);
  uint64 total_w = end_w - begin_w;
  if (total_w == 0) return NULL;
  if (strategy == SAMPLE_UNIFORM_OVER_ADS) 
    return begin_p + rng->Random() % (end_p - begin_p);

  uint64 target_weight = begin_w + (rng->Random() % total_w) + 1;
  Tuple weight_tuple; weight_tuple.weight_ = target_weight;
  return lower_bound(begin_p, end_p, weight_tuple, WeightCompare());
}

uint64 Sampler::RangeWeight(const Tuple & begin, const Tuple & end) {
  CHECK(built_);
  Tuple * begin_p = lower_bound(data_, data_ + size_, begin);
  Tuple * end_p = lower_bound(data_, data_ + size_, end);
  uint64 begin_w = WeightBefore(begin_p);
  uint64 end_w = WeightBefore(end_p);
  //cout << "RangeWeight " << begin.ToString() << " " << end.ToString() 
  //	 << " begin_p = " << (begin_p-data_) << " end_p = " << (end_p-data_) 
  //	 << " begin_w = " << begin_w << " end_w = " << end_w << endl;
  CHECK(end_w >= begin_w);
  uint64 total_w = end_w - begin_w;
  return total_w;
}

uint64 Sampler::RangeSize(const Tuple & begin, const Tuple & end) {
  CHECK(built_);
  Tuple * begin_p = lower_bound(data_, data_ + size_, begin);
  Tuple * end_p = lower_bound(data_, data_ + size_, end);
  return end_p - begin_p;
}

void Sampler::CopyOutRangeLocked(const Tuple & begin, const Tuple &end,
				 uint64 **destination, uint64 * source_end_p) {
  Tuple * begin_p = lower_bound(data_, data_ + size_, begin);
  Tuple * end_p = lower_bound(data_, data_ + size_, end);
  for (Tuple *run = begin_p; run < end_p; run++) {
    if (*destination == source_end_p) break;
    // cout << "lhs: " << ::ToString(Thing(run->a_[LHS])) << endl;
    // cout << "rhs: " << ::ToString(Thing(run->a_[RHS])) << endl;
    // cout << "ad: " << ::ToString(Thing(run->a_[AD])) << endl;
    **destination = run->a_[AD];
    Thing adcheck(**destination); CHECK(adcheck.GetBox());
    (*destination)++;
    uint64 w = run->weight_ - WeightBefore(run);
    // cout << "w: " << w << endl;
    **destination = Thing::MakeInteger(w).ID();
    (*destination)++;
    // cout << "*destination = " << (uint64) (*destination) << endl;
  }
}

uint64 Sampler::WeightPerNode(uint64 k, uint64 n) {
  if (size_ == 0) return 1;
  RNG rng;
  vector<uint64> wpn;
  for (uint64 i=0; i<n; i++) {
    Tuple &t = data_[rng.Random() % size_];
    if (false)
      cout << i << " w:" 
	   << t.weight_
	   << " nc:" << t.NodeCount()
	   << endl;
    wpn.push_back(t.WeightPerNode());
  }
  sort(wpn.begin(), wpn.end());
  return wpn[k];
}

void Sampler::Test() {
  Sampler s;
  s.Allocate(1000);
  ifstream input("test.data");
  Tuple r;
  RNG rng;
  while (r.Read(input)) {
    cout << "Adding tuple " << r.ToString() << endl;
    s.AddTuple(r);
  }
  s.Build();
  cout << s.ToString() << endl;
  string cmd;
  while (cin >> cmd) {
    if (cmd == "w") {
      r.ReadPrefix(cin);
      cout << "Range weight " << r.ToString() 
	   << " = " << s.RangeWeight(r.PrefixToLowerBound(), r.PrefixToUpperBound()) 
	   << endl;
    }
    if (cmd == "s") {
      r.ReadPrefix(cin);
      uint64 reps;
      cin >> reps;
      for (uint64 rep=0; rep<reps; rep++) {
	Tuple * sample = s.Sample(r.PrefixToLowerBound(),
				  r.PrefixToUpperBound(), &rng, SAMPLE_WEIGHTED);
	cout << r.ToString() << " -> " 
	     << (sample ? sample->ToString() : string("MISSING")) << endl;
      }
    }
  }
}
    
AdStore::AdStore() :
  samplers_(1000), trash_tuples_(0), used_tuples_(0) {
  sample_lock_ = new RWLock;
  sample_lock_reader_ = new RWLock::Reader(sample_lock_);
  pthread_mutex_init(&modify_lock_, NULL);
  max_tuples_ = 1 << ConfigInt("lg_adstore_max_tuples");
}

Tuple *AdStore::Sample
(const Tuple & prefix, RWLock::Reader *reader, RNG *rng, SampleStrategy strategy) {
  return Sample(prefix.PrefixToLowerBound(), 
		prefix.PrefixToUpperBound(), reader, rng, strategy); }

Tuple *AdStore::Sample(const Tuple & begin, const Tuple & end, 
		       RWLock::Reader *reader, 
		       RNG *rng, SampleStrategy strategy) {
  RWLock::ReadLocker rlocker(reader);
  uint64 sampler_weights[100];
  CHECK(samplers_.size() <= 100);
  uint64 sum_sampler_weights = 0;
  for (uint64 i=0; i<samplers_.size(); i++) {
    Sampler * s = samplers_[i];
    uint64 w = 0;
    switch(strategy) {
    case SAMPLE_WEIGHTED:
      w = s->RangeWeight(begin, end); break;
    case SAMPLE_UNIFORM_OVER_SAMPLERS: 
      w = (s->RangeSize(begin, end) ? 1 : 0);
      break;
    case SAMPLE_UNIFORM_OVER_ADS:
      w = s->RangeSize(begin, end); break;
    default: CHECK(false); break;
    }
    sampler_weights[i] = w;
    sum_sampler_weights += w;
  }
  if (sum_sampler_weights == 0) return NULL;
  uint64 r = rng->Random() % sum_sampler_weights;
  uint64 sum = 0;
  for (uint64 i=0; i<samplers_.size(); i++) {
    sum += sampler_weights[i];
    if (sum > r) return samplers_[i]->Sample(begin, end, rng, strategy);	
  }
  return NULL;
}
  
uint64 AdStore::RangeWeight(const Tuple & prefix, RWLock::Reader *reader) {
  return RangeWeight(prefix.PrefixToLowerBound(), 
		     prefix.PrefixToUpperBound(), reader); }

uint64 AdStore::RangeWeight(const Tuple & begin, const Tuple & end,
			    RWLock::Reader *reader) {
  RWLock::ReadLocker rlocker(reader);
  uint64 ret = 0;
  for (uint64 i=0; i<samplers_.size(); i++) 
    ret += samplers_[i]->RangeWeight(begin, end);
  return ret;
}

// assume that we have a read lock before calling this.
uint64 AdStore::RangeSizeLocked(const Tuple & begin, const Tuple &end) {
  uint64 ret = 0;
  for (uint64 i=0; i<samplers_.size(); i++) 
    ret += samplers_[i]->RangeSize(begin, end);
  return ret;
}

// assumes we have a read lock 
// copies out the AD field to the destination
void AdStore::CopyOutRangeLocked(const Tuple &begin, 
				 const Tuple &end, 
				 uint64 **destination,
				 uint64 * source_end_p) {
  for (uint64 i=0; i<samplers_.size(); i++) 
    samplers_[i]->
      CopyOutRangeLocked(begin, end, destination, source_end_p);
}

void AdStore::AddSamplerI(Sampler *s) {
  // cout << "::AddSamplerI" << endl;
  samplers_.push_back(s);
  used_tuples_ += s->size_;
  while (NeedsMerging()) MergeTopTwoSamplers();
  while (TooManyTuples()) {
    RWLock::WriteLocker locker(sample_lock_reader_);    
    if (TooManyTuples()) Shrink();
  }
}

bool AdStore::AddSamplers(Sampler **s, uint64 num_samplers, bool force) {
  if (force) pthread_mutex_lock(&modify_lock_);
  else {
    int result = pthread_mutex_trylock(&modify_lock_);
    if (result) return false;
  }
  for (uint64 i=0; i<num_samplers; i++) AddSamplerI(s[i]);
  pthread_mutex_unlock(&modify_lock_);
  return true;
}

bool AdStore::NeedsMerging() {
  return ( (samplers_.size() > num_permanent_ + 1) 
	   && (samplers_[samplers_.size()-1]->size_ >= 
	       (samplers_[samplers_.size()-2]->size_ / merge_multiple_)));
}

// Merge all samplers
void AdStore::MergeAllNonPermanentI() {
  while (samplers_.size() > num_permanent_ + 1) 
    MergeTopTwoSamplers(); 
  if (samplers_.size() == num_permanent_)
    samplers_.push_back(MakeEmptySampler());
}

// Merge two samplers
void AdStore::MergeTopTwoSamplers() {
  Sampler *as1 = samplers_[samplers_.size()-1];
  Sampler *as2 = samplers_[samplers_.size()-2];
  trash_tuples_ += as1->size_ + as2->size_;
  if (TooMuchTrash()) EmptyTrash();
  Sampler *as = Sampler::Merge(as1, as2);
  samplers_.pop_back();
  samplers_.back() = as;
  trash_.push_back(as1);
  trash_.push_back(as2);
}

bool AdStore::TooManyTuples() {
  return (used_tuples_ * 2 > max_tuples_);
}

bool AdStore::TooMuchTrash() {
  return (used_tuples_ + trash_tuples_ > max_tuples_);
}
  
void AdStore::EmptyTrash() {
  if (trash_.size() == 0) return;
  RWLock::WriteLocker locker(sample_lock_reader_);
  for (uint64 i=0; i<trash_.size(); i++) {
    Sampler *as = trash_[i];
    trash_tuples_ -= as->size_;
    delete as;
  }
  trash_.clear();
}

void AdStore::MakePermanent() {
  MutexLocker locker(&modify_lock_);
  MergeAllNonPermanentI();
  num_permanent_ = samplers_.size(); 
}

void AdStore::Store(string filename) {
  MergeAllNonPermanentI();
  OnlyNonpermanentSampler()->Store(filename);
}

// Assumes permanent ads have been loaded already
void AdStore::Load(string filename) {
  Sampler * s = new Sampler;
  s->Load(filename);
  samplers_.push_back(s);
  used_tuples_ += s->size_;
}

struct JudgeTask : public JobQueue::Task {
  void Run() {
    uint64 local_num_keep = 0;
    RNG rng;
    for (Ads::Tuple * run = start_; run<end_; run++) {
      uint64 min_weight = judge_->MinWeight(*run);
      if (run->weight_ < min_weight)
	run->weight_ = 
	  ((rng.Random() % min_weight) < run->weight_) ? min_weight : 0;
      if (run->weight_) local_num_keep++;
    }
    __sync_fetch_and_add(num_keep_, local_num_keep);
  }
  
  JudgeTask(Ads::Tuple * start, Ads::Tuple * end, 
	    TupleJudge * judge,
  	    uint64 * num_keep)
    :start_(start), end_(end), judge_(judge), num_keep_(num_keep) {}

  Ads::Tuple * start_;
  Ads::Tuple * end_;
  TupleJudge * judge_;
  uint64 * num_keep_;
};

struct SimpleTupleJudge : public Ads::TupleJudge{
  SimpleTupleJudge(uint64 mwpn) :min_weight_per_node_(mwpn) {}
  uint64 MinWeight(const Ads::Tuple & t) {
    uint64 node_count = 0;
    for (uint64 i=0; i<Ads::WIDTH; i++) 
      node_count += Thing(t.a_[i]).NodeCount();
    return min_weight_per_node_ * node_count;
  }
  uint64 min_weight_per_node_;
};

void AdStore::Shrink() {
  MergeAllNonPermanentI();
  if (samplers_.size() != num_permanent_ + 1) {
    cout << "Unexpected number of samplers" << samplers_.size() 
	 << " num_permanent_ = " << num_permanent_;
    return;
  }
  Sampler * s = samplers_.back();
  s->UnSumWeights();
  uint64 median_weight = 2 * s->WeightPerNode(500, 1000);
  cout << "::AdStore::Shrink median_weight:" << median_weight << endl;
  SimpleTupleJudge judge(median_weight);
  uint64 num_keep = 0;
  JobQueue::Job * j = new JobQueue::Job(G_JQ);
  uint64 num_shards = 100;
  for (uint64 i=0; i<num_shards; i++) {
    j->AddTask(new JudgeTask(s->data_ + ShardBegin(s->size_, num_shards, i), 
			     s->data_ + ShardEnd(s->size_, num_shards, i),
			     &judge, &num_keep));    
  }
  j->Wait();
  delete j;
  Sampler * new_s = new Sampler;
  new_s->Allocate(num_keep);
  for (uint64 i=0; i<s->size_; i++) {
    if (s->data_[i].weight_) 
      new_s->AddTuple(s->data_[i]);    
  }
  new_s->SumWeights();
  new_s->built_ = true;
  new_s->Verify();
  used_tuples_ -= (s->size_ - new_s->size_);
  samplers_.back() = new_s;
  delete s;
}

AdStore * TEST_ADSTORE;

// Create some random ads and add to the sampler
const uint64 test_addads_range = 10;
const uint64 test_addads_range_weight = 10;
const uint64 test_addads_range_number = 1000;

void TestAddAds() {
  Tuple t;
  RNG r;
  Sampler * s = new Sampler();
  uint64 number = r.Random() % test_addads_range_number;
  s->Allocate(number);
  for (uint64 i=0; i<number; i++) {
    for (uint64 j=0; j < WIDTH; j++)
      t.a_[j] = r.Random() % test_addads_range;
    t.weight_ = r.Random() % test_addads_range_weight;
    s->AddTuple(t);
  }
  s->Build();
  COUT << "About to add sampler " << number << " tuples\n";
  TEST_ADSTORE->AddSamplers(&s, 1, true);
  COUT << "Added sampler " << number << " tuples\n";
}

void * LoopAddAds(void *) {
  while (1) TestAddAds(); return NULL;
}

void TestReadAds(RWLock::Reader *reader) {
  RNG rng;
  uint64 prefix = rng.Random() % test_addads_range;
  Tuple pt = Ads::Prefix(1, prefix);
  Tuple * sampled = TEST_ADSTORE->Sample(pt, reader, &rng, SAMPLE_WEIGHTED);
  COUT << pt.ToString() << " -> " << (sampled ? sampled->ToString() 
				      : string("absent"))<< "\n";
}

void * LoopReadAds(void *) {
  RWLock::Reader reader(TEST_ADSTORE->sample_lock_);
  while (1) 
    TestReadAds(&reader); return NULL;
}

const uint64 test_num_write_threads = 5;
const uint64 test_num_read_threads = 5;
void AdStore::Test() {
  TEST_ADSTORE = new AdStore();
  for (uint64 i=0; i<test_num_write_threads; i++) {
    pthread_t * thread = new pthread_t;
    pthread_create(thread, NULL, LoopAddAds, NULL);
  }
  for (uint64 i=0; i<test_num_read_threads; i++) {
    pthread_t * thread = new pthread_t;
    pthread_create(thread, NULL, LoopReadAds, NULL);
  }
  while(1) { sleep (10); }
}

TupleJudge::~TupleJudge() {}


CLOSEBRACE

  
int test_ads_main() { 
  // Ads::Sampler::Test();  
  InitUtil(0, NULL);
  Ads::AdStore::Test();
  return 0; }


